# app/routes.py

# 导入所需的库和模块
from flask import Blueprint, request, jsonify, session, redirect, url_for, flash, render_template
import os  # 用于处理文件路径
from config import Config  # 导入配置文件
from src.image_processing.image_analysis import analyze_image  # 导入图像分析函数
from src.image_processing.key_sensitivity_analysis import analyze_key_sensitivity  # 导入密钥敏感性分析函数
from src.image_processing.image_handler import ImageData, save_file_async, generate_unique_filename  # 导入图像处理类和相关函数
from concurrent.futures import ThreadPoolExecutor  # 用于并发执行任务


# 创建一个名为 'main' 的蓝图
main_bp = Blueprint('main', __name__)


# 定义根路由，支持 GET 和 POST 方法
@main_bp.route('/', methods=['GET', 'POST'])
async def index():
    if request.method == 'POST':  # 如果是 POST 请求
        if 'file' not in request.files:  # 检查是否有文件上传
            return jsonify({'success': False, 'message': '未选择文件'})  # 返回错误信息

        file = request.files['file']  # 获取上传的文件
        if file.filename == '':  # 检查文件名是否为空
            return jsonify({'success': False, 'message': '未选择文件'})  # 返回错误信息

        if file:  # 如果有文件上传
            file_extension = file.filename.split('.')[-1].lower()  # 获取文件扩展名并转换为小写
            if file_extension not in Config.ALLOWED_EXTENSIONS:  # 检查文件类型是否允许
                return jsonify({'success': False, 'message': '不支持的文件类型，请上传PNG、JPG或JPEG格式的图像。'})  # 返回错误信息

            # 生成唯一的文件名
            filename = generate_unique_filename(file.filename)
            filepath = os.path.join('app/static/uploads', filename)  # 构建文件保存路径
            await save_file_async(file, filepath)  # 异步保存文件

            # 使用 ImageData 类加载图像
            image_data = ImageData(image_path=filepath)
            mode = image_data.get_mode()  # 获取图像的模式（如 RGB、L 等）

            # 返回成功信息和图像类型
            return jsonify({
                'success': True,
                'image_type': f'{file_extension.upper()} ({mode})',
                'original_mode': mode,
                'converted_mode': mode
            })

    # 如果是 GET 请求，渲染 index.html 模板
    return render_template('index.html')

# 定义加密路由，仅支持 POST 方法
@main_bp.route('/encrypt', methods=['POST'])
async def encrypt():
    if 'file' not in request.files:  # 检查是否有文件上传
        return jsonify({'success': False, 'message': '未选择文件'})  # 返回错误信息

    file = request.files['file']  # 获取上传的文件
    if file.filename == '':  # 检查文件名是否为空
        return jsonify({'success': False, 'message': '未选择文件'})  # 返回错误信息

    if file:  # 如果有文件上传
        file_extension = file.filename.split('.')[-1].lower()  # 获取文件扩展名并转换为小写
        if file_extension not in Config.ALLOWED_EXTENSIONS:  # 检查文件类型是否允许
            return jsonify({'success': False, 'message': '不支持的文件类型，请上传PNG、JPG或JPEG格式的图像。'})  # 返回错误信息

        # 生成唯一的文件名
        filename = generate_unique_filename(file.filename)
        filepath = os.path.join('app/static/uploads', filename)  # 构建文件保存路径
        await save_file_async(file, filepath)  # 异步保存文件

        key = request.form.get('key')  # 获取表单中的密钥
        if not all(c in '01' for c in key):  # 检查密钥是否为二进制字符串（仅包含0和1）
            return jsonify({'success': False, 'message': '密钥必须为二进制字符串（仅包含0和1）'})  # 返回错误信息

        key = key.zfill(21)  # 将密钥填充到21位

        # 使用 ImageData 类加载图像
        image_data = ImageData(image_path=filepath)

        try:
            # 使用线程池并行处理 RC4 和 LFSR 加密和解密
            with ThreadPoolExecutor() as executor:
                # 加密
                rc4_encrypted_future = executor.submit(
                    image_data.encrypt, key, 'rc4', os.path.join('app/static/uploads', 'rc4_encrypted_'+ filename)
                )
                lfsr_encrypted_future = executor.submit(
                    image_data.encrypt, key, 'lfsr', os.path.join('app/static/uploads', 'lfsr_encrypted_'+ filename)
                )

                rc4_encrypted_data = rc4_encrypted_future.result()  # 获取 RC4 加密结果
                lfsr_encrypted_data = lfsr_encrypted_future.result()  # 获取 LFSR 加密结果

                # 解密
                executor.submit(
                    rc4_encrypted_data.encrypt, key, 'rc4', os.path.join('app/static/uploads', 'rc4_decrypted_' + filename)
                )
                executor.submit(
                    lfsr_encrypted_data.encrypt, key, 'lfsr', os.path.join('app/static/uploads', 'lfsr_decrypted_'+ filename)
                )

        except Exception as e:  # 捕获异常
            return jsonify({'success': False, 'message': f'加密或解密失败: {str(e)}'})  # 返回错误信息

        # 分析原始图像
        original_analysis = analyze_image(image_data, 'Original Image_'+filename, 'app/static/uploads')

        # 分析 RC4 加密图像
        rc4_encrypted_analysis = analyze_image(rc4_encrypted_data, 'RC4 Encrypted Image_'+filename, 'app/static/uploads', image_data)

        # 分析 LFSR 加密图像
        lfsr_encrypted_analysis = analyze_image(lfsr_encrypted_data, 'LFSR Encrypted Image_'+filename, 'app/static/uploads', image_data)

        # 修改密钥以进行密钥敏感性分析
        modified_key = key[:-1] + ('1' if key[-1] == '0' else '0')
        if all(c == '0' for c in modified_key):
            modified_key = '1' + modified_key[1:]

        # LFSR 密钥敏感性分析
        lfsr_npcr_values, lfsr_uaci_values, image_mode = analyze_key_sensitivity(image_data, key, modified_key, 'lfsr', 'app/static/uploads','lfsr_'+filename)
        # RC4 密钥敏感性分析
        rc4_npcr_values, rc4_uaci_values, _ = analyze_key_sensitivity(image_data, key, modified_key, 'rc4', 'app/static/uploads','rc4_'+filename)

        # 确定通道名称
        if image_mode == 'L':
            channel_names = ['Gray']
        elif image_mode == 'RGB':
            channel_names = ['Red', 'Green', 'Blue']
        elif image_mode == 'RGBA':
            channel_names = ['Red', 'Green', 'Blue', 'Alpha']
        else:
            channel_names = [f'Channel {i + 1}' for i in range(len(rc4_npcr_values))]

        result_data = {
            'original_image': filename,
            'rc4_encrypted_image': 'rc4_encrypted_'+filename,
            'lfsr_encrypted_image': 'lfsr_encrypted_'+filename,
            'rc4_decrypted_image': 'rc4_decrypted_'+filename,
            'lfsr_decrypted_image': 'lfsr_decrypted_'+filename,
            'original_analysis': original_analysis,
            'rc4_encrypted_analysis': rc4_encrypted_analysis,
            'lfsr_encrypted_analysis': lfsr_encrypted_analysis,
            'rc4_npcr_values': rc4_npcr_values,
            'rc4_uaci_values': rc4_uaci_values,
            'lfsr_npcr_values': lfsr_npcr_values,
            'lfsr_uaci_values': lfsr_uaci_values,
            'channel_names': channel_names,
            'rc4_psnr': rc4_encrypted_analysis.get('psnr', [None] * len(channel_names)),
            'lfsr_psnr': lfsr_encrypted_analysis.get('psnr', [None] * len(channel_names)),
            'rc4_chi_square': rc4_encrypted_analysis.get('chi_square', [None] * len(channel_names)),
            'lfsr_chi_square': lfsr_encrypted_analysis.get('chi_square', [None] * len(channel_names))
        }

        # 将结果存储在 session 中
        session['result'] = result_data

        # 返回成功信息
        return jsonify({'success': True})

@main_bp.route('/result')
def result():
    result_data = session.get('result')  # 从 session 中获取结果数据
    if not result_data:  # 如果没有结果数据
        flash('未找到加密结果，请先上传并加密图像。')  # 显示提示信息
        return redirect(url_for('main.index'))  # 重定向到首页

    # 确保只处理一次结果
    if 'processed' not in session:
        session['processed'] = True

    # 渲染 result.html 模板，并传递结果数据
    return render_template('result.html', result_data=result_data)
